import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import yaml
import os
import threading
import subprocess
import sys
from pathlib import Path
import time
import json


class YOLOTrainingGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("YOLO模型训练可视化界面")
        self.root.geometry("1400x900")
        self.root.configure(bg='white')

        # 训练进程控制
        self.training_process = None
        self.training_thread = None
        self.is_training = False

        # 界面样式配置
        self.setup_styles()
        self.create_widgets()
        self.setup_tooltips()

        # 延迟创建菜单栏以避免循环导入
        self.root.after(100, self.create_menu_bar)

    def create_menu_bar(self):
        """延迟创建菜单栏"""
        try:
            from UI.ui.menu import MenuBar
            self.menu_bar = MenuBar(self.root, None, self)  # 传入self作为trainer参数
        except Exception as e:
            print(f"创建菜单栏失败: {e}")


    def setup_styles(self):
        """设置界面样式"""
        style = ttk.Style()
        style.theme_use('clam')  # 使用clam主题
        
        # 配置颜色方案
        style.configure('TFrame', background='white')
        style.configure('TLabel', background='white', foreground='black')
        style.configure('TButton', background='#4CAF50', foreground='white', font=('微软雅黑', 10, 'bold'))
        style.map('TButton', 
                 background=[('active', '#45a049')],
                 foreground=[('active', 'white')])
        
        # 按钮样式
        style.configure('Start.TButton', background='#4CAF50', foreground='white', font=('微软雅黑', 10, 'bold'))
        style.map('Start.TButton', 
                 background=[('active', '#45a049')],
                 foreground=[('active', 'white')])
        
        style.configure('Stop.TButton', background='#f44336', foreground='white', font=('微软雅黑', 10, 'bold'))
        style.map('Stop.TButton', 
                 background=[('active', '#d32f2f')],
                 foreground=[('active', 'white')])
        
        # 树形控件样式
        style.configure('Treeview', 
                       background='white', 
                       foreground='black',
                       fieldbackground='white',
                       font=('微软雅黑', 9))
        style.configure('Treeview.Heading', 
                       background='#e0e0e0', 
                       foreground='black',
                       font=('微软雅黑', 10, 'bold'))
        style.map('Treeview', 
                 background=[('selected', '#b8d1f3')])
        
        # 文本框样式
        style.configure('Log.TText', 
                       background='black', 
                       foreground='white',
                       font=('Consolas', 10))
        
    def create_widgets(self):
        # 主框架
        main_frame = ttk.Frame(self.root, padding="15")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 设置网格权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(3, weight=1)
        
        # 左侧主面板
        left_panel = ttk.Frame(main_frame)
        left_panel.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S), padx=(0, 10), pady=(0, 10))
        
        # 权重文件选择
        weight_frame = ttk.LabelFrame(left_panel, text="权重文件", padding="10", relief=tk.RAISED)
        weight_frame.grid(row=0, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Label(weight_frame, text="路径:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        self.weight_path_var = tk.StringVar()
        self.weight_path_entry = ttk.Entry(weight_frame, textvariable=self.weight_path_var, width=60)
        self.weight_path_entry.grid(row=0, column=1, padx=(0, 5), sticky=(tk.W, tk.E))
        ttk.Button(weight_frame, text="浏览", command=self.browse_weight_file).grid(row=0, column=2, padx=(0, 5))
        
        # 数据集配置
        dataset_frame = ttk.LabelFrame(left_panel, text="数据集配置", padding="10", relief=tk.RAISED)
        dataset_frame.grid(row=1, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Label(dataset_frame, text="data.yaml:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        self.data_yaml_var = tk.StringVar()
        self.data_yaml_entry = ttk.Entry(dataset_frame, textvariable=self.data_yaml_var, width=60)
        self.data_yaml_entry.grid(row=0, column=1, padx=(0, 5), sticky=(tk.W, tk.E))
        ttk.Button(dataset_frame, text="浏览", command=self.browse_data_yaml).grid(row=0, column=2, padx=(0, 5))
        
        # 训练配置
        config_frame = ttk.LabelFrame(left_panel, text="训练参数", padding="10", relief=tk.RAISED)
        config_frame.grid(row=2, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Label(config_frame, text="train.yaml:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        self.train_yaml_var = tk.StringVar()
        self.train_yaml_entry = ttk.Entry(config_frame, textvariable=self.train_yaml_var, width=60)
        self.train_yaml_entry.grid(row=0, column=1, padx=(0, 5), sticky=(tk.W, tk.E))
        ttk.Button(config_frame, text="浏览", command=self.browse_train_yaml).grid(row=0, column=2, padx=(0, 5))
        
        # 参数显示区域
        params_frame = ttk.LabelFrame(left_panel, text="参数详情", padding="10", relief=tk.RAISED)
        params_frame.grid(row=3, column=0, sticky=(tk.W, tk.E, tk.N, tk.S), pady=(0, 10))
        
        # 创建参数表格
        self.params_tree = ttk.Treeview(params_frame, columns=('参数名', '值'), show='tree headings', height=8)
        self.params_tree.heading('参数名', text='参数名', anchor=tk.W)
        self.params_tree.heading('值', text='值', anchor=tk.W)
        self.params_tree.column('参数名', width=200, minwidth=100, anchor=tk.W)
        self.params_tree.column('值', width=300, minwidth=100, anchor=tk.W)
        
        # 添加滚动条
        params_scrollbar = ttk.Scrollbar(params_frame, orient=tk.VERTICAL, command=self.params_tree.yview)
        self.params_tree.configure(yscrollcommand=params_scrollbar.set)
        
        self.params_tree.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        params_scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
        
        # 保存路径选择
        save_frame = ttk.LabelFrame(left_panel, text="结果保存", padding="10", relief=tk.RAISED)
        save_frame.grid(row=4, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Label(save_frame, text="路径:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        self.save_path_var = tk.StringVar()
        self.save_path_entry = ttk.Entry(save_frame, textvariable=self.save_path_var, width=60)
        self.save_path_entry.grid(row=0, column=1, padx=(0, 5), sticky=(tk.W, tk.E))
        ttk.Button(save_frame, text="浏览", command=self.browse_save_path).grid(row=0, column=2, padx=(0, 5))
        
        # 控制按钮
        control_frame = ttk.Frame(left_panel)
        control_frame.grid(row=5, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        self.start_button = ttk.Button(control_frame, text="开始训练", command=self.start_training, style='Start.TButton')
        self.start_button.grid(row=0, column=0, padx=(0, 10))
        
        self.stop_button = ttk.Button(control_frame, text="停止训练", command=self.stop_training, state='disabled', style='Stop.TButton')
        self.stop_button.grid(row=0, column=1, padx=(0, 10))
        
        self.load_config_button = ttk.Button(control_frame, text="加载配置", command=self.load_configurations)
        self.load_config_button.grid(row=0, column=2, padx=(0, 10))
        
        # 右侧日志面板
        log_frame = ttk.LabelFrame(main_frame, text="训练日志", padding="10", relief=tk.RAISED)
        log_frame.grid(row=0, column=1, sticky=(tk.W, tk.E, tk.N, tk.S), padx=(10, 0), pady=(0, 10))
        
        # 创建日志文本框
        self.log_text = tk.Text(log_frame, bg='black', fg='white', font=('Consolas', 10), wrap=tk.NONE)
        log_scrollbar = ttk.Scrollbar(log_frame, orient=tk.VERTICAL, command=self.log_text.yview)
        self.log_text.configure(yscrollcommand=log_scrollbar.set)
        
        self.log_text.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        log_scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
        
        # 配置网格权重
        left_panel.columnconfigure(0, weight=1)
        left_panel.rowconfigure(3, weight=1)
        params_frame.columnconfigure(0, weight=1)
        params_frame.rowconfigure(0, weight=1)
        log_frame.columnconfigure(0, weight=1)
        log_frame.rowconfigure(0, weight=1)
        
    def setup_tooltips(self):
        """设置鼠标悬停提示"""
        self.tooltip_texts = {
            'weights': '预训练权重文件路径，用于初始化模型参数',
            'data': '数据集配置文件路径，包含训练/验证数据路径和类别信息',
            'epochs': '训练轮数，模型遍历整个数据集的次数',
            'batch_size': '批处理大小，每次训练使用的样本数量',
            'img_size': '输入图像尺寸，通常为32的倍数',
            'learning_rate': '学习率，控制模型参数更新的步长',
            'save_path': '训练结果保存路径，包含模型权重和日志文件'
        }
        
    def browse_weight_file(self):
        filename = filedialog.askopenfilename(
            title="选择权重文件",
            filetypes=[("权重文件", "*.pt"), ("所有文件", "*.*")]
        )
        if filename:
            self.weight_path_var.set(filename)
            
    def browse_data_yaml(self):
        filename = filedialog.askopenfilename(
            title="选择data.yaml文件",
            filetypes=[("YAML文件", "*.yaml"), ("YAML文件", "*.yml"), ("所有文件", "*.*")]
        )
        if filename:
            self.data_yaml_var.set(filename)
            
    def browse_train_yaml(self):
        filename = filedialog.askopenfilename(
            title="选择train.yaml文件",
            filetypes=[("YAML文件", "*.yaml"), ("YAML文件", "*.yml"), ("所有文件", "*.*")]
        )
        if filename:
            self.train_yaml_var.set(filename)
            
    def browse_save_path(self):
        directory = filedialog.askdirectory(title="选择保存路径")
        if directory:
            self.save_path_var.set(directory)
            
    def load_configurations(self):
        """加载配置文件并填充参数"""
        try:
            # 清空现有参数
            for item in self.params_tree.get_children():
                self.params_tree.delete(item)
            
            # 加载data.yaml
            data_yaml_path = self.data_yaml_var.get()
            if data_yaml_path and os.path.exists(data_yaml_path):
                with open(data_yaml_path, 'r', encoding='utf-8') as f:
                    data_config = yaml.safe_load(f)
                
                self.params_tree.insert('', 'end', values=('数据集路径', data_config.get('path', '')))
                self.params_tree.insert('', 'end', values=('训练集', data_config.get('train', '')))
                self.params_tree.insert('', 'end', values=('验证集', data_config.get('val', '')))
                self.params_tree.insert('', 'end', values=('类别数', str(data_config.get('nc', ''))))
                if 'names' in data_config:
                    names_str = str(data_config['names'])
                    if len(names_str) > 50:
                        names_str = names_str[:50] + '...'
                    self.params_tree.insert('', 'end', values=('类别名称', names_str))
            
            # 加载train.yaml
            train_yaml_path = self.train_yaml_var.get()
            if train_yaml_path and os.path.exists(train_yaml_path):
                with open(train_yaml_path, 'r', encoding='utf-8') as f:
                    train_config = yaml.safe_load(f)
                
                for key, value in train_config.items():
                    self.params_tree.insert('', 'end', values=(key, str(value)))
                    
        except Exception as e:
            messagebox.showerror("错误", f"加载配置文件失败: {str(e)}")
            
    def start_training(self):
        """开始训练"""
        if not self.validate_inputs():
            return
            
        self.is_training = True
        self.start_button.config(state='disabled')
        self.stop_button.config(state='normal')
        self.log_text.delete(1.0, tk.END)
        self.log_text.insert(tk.END, "开始训练...\n")
        
        # 在新线程中运行训练
        self.training_thread = threading.Thread(target=self.run_training)
        self.training_thread.daemon = True
        self.training_thread.start()
        
    def stop_training(self):
        """停止训练"""
        if self.training_process:
            self.training_process.terminate()
        self.is_training = False
        self.start_button.config(state='normal')
        self.stop_button.config(state='disabled')
        self.log_text.insert(tk.END, "\n训练已停止\n")
        
    def validate_inputs(self):
        """验证输入"""
        if not self.weight_path_var.get():
            messagebox.showerror("错误", "请选择权重文件")
            return False
        if not self.data_yaml_var.get():
            messagebox.showerror("错误", "请选择data.yaml文件")
            return False
        if not self.train_yaml_var.get():
            messagebox.showerror("错误", "请选择train.yaml文件")
            return False
        if not self.save_path_var.get():
            messagebox.showerror("错误", "请选择保存路径")
            return False
            
        # 检查文件是否存在
        if not os.path.exists(self.weight_path_var.get()):
            messagebox.showerror("错误", "权重文件不存在")
            return False
        if not os.path.exists(self.data_yaml_var.get()):
            messagebox.showerror("错误", "data.yaml文件不存在")
            return False
        if not os.path.exists(self.train_yaml_var.get()):
            messagebox.showerror("错误", "train.yaml文件不存在")
            return False
            
        return True
        
    def run_training(self):
        """运行实际的训练过程"""
        try:
            # 创建临时训练脚本文件
            script_content = self.generate_training_script()
            script_path = "temp_train_script.py"
            
            with open(script_path, 'w', encoding='utf-8') as f:
                f.write(script_content)
            
            # 构建训练命令
            cmd = [sys.executable, script_path]
            
            # 执行训练
            self.training_process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                universal_newlines=True,
                bufsize=1,
                encoding='utf-8'
            )
            
            # 实时读取输出
            for line in iter(self.training_process.stdout.readline, ''):
                if not self.is_training:
                    break
                if line.strip():  # 只显示非空行
                    self.log_text.insert(tk.END, line)
                    self.log_text.see(tk.END)
                    self.root.update_idletasks()
                
            self.training_process.wait()
            
            # 清理临时文件
            if os.path.exists(script_path):
                os.remove(script_path)
                
        except Exception as e:
            self.log_text.insert(tk.END, f"训练出错: {str(e)}\n")
        finally:
            self.is_training = False
            self.start_button.config(state='normal')
            self.stop_button.config(state='disabled')
            self.log_text.insert(tk.END, "\n训练完成\n")
            
    def generate_training_script(self):
        """生成训练脚本"""
        weight_path = self.weight_path_var.get().replace('\\', '/')
        data_yaml_path = self.data_yaml_var.get().replace('\\', '/')
        train_yaml_path = self.train_yaml_var.get().replace('\\', '/')
        save_path = self.save_path_var.get().replace('\\', '/')
        
        # 读取训练配置并转换为Python字典格式
        try:
            with open(train_yaml_path, 'r', encoding='utf-8') as f:
                train_config = yaml.safe_load(f)
        except Exception as e:
            train_config = {}
            self.log_text.insert(tk.END, f"读取训练配置失败: {str(e)}\n")
            
        # 将配置转换为Python字典的字符串表示
        config_str = self.dict_to_python_str(train_config)
        
        # 构建训练脚本
        script = f'''# -*- coding: utf-8 -*-
import sys
import os
import time

def train_model():
    try:
        # 导入ultralytics
        try:
            from ultralytics import YOLO
            print("成功导入ultralytics")
        except ImportError as e:
            print(f"导入ultralytics失败: {{e}}")
            print("请安装: pip install ultralytics")
            return
            
        # 加载模型
        print("加载模型: {weight_path}")
        try:
            model = YOLO('{weight_path}')
            print("模型加载成功")
        except Exception as e:
            print(f"模型加载失败: {{e}}")
            return
        
        # 训练配置
        train_config = {config_str}
        train_config['data'] = r'{data_yaml_path}'
        train_config['project'] = r'{save_path}'
        train_config['name'] = 'yolo_training'
        
        print("训练配置:")
        for key, value in train_config.items():
            print(f"  {{key}}: {{value}}")
        
        print("开始训练...")
        # 开始训练
        results = model.train(**train_config)
        
        print("训练完成!")
        print("结果保存在:", results)
        
    except Exception as e:
        print(f"训练出错: {{e}}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    train_model()
'''
        return script
        
    def dict_to_python_str(self, d):
        """将字典转换为Python字典字符串，正确处理布尔值和None"""
        if isinstance(d, dict):
            items = []
            for k, v in d.items():
                key = repr(k)
                value = self.value_to_python_str(v)
                items.append(f"{key}: {value}")
            return "{" + ", ".join(items) + "}"
        return repr(d)
        
    def value_to_python_str(self, value):
        """将值转换为Python字符串表示"""
        if isinstance(value, bool):
            return "True" if value else "False"
        elif value is None or (isinstance(value, str) and value.lower() == "none"):
            return "None"
        elif isinstance(value, (int, float)):
            return repr(value)
        elif isinstance(value, str):
            # 检查是否是布尔值字符串
            if value.lower() == "true":
                return "True"
            elif value.lower() == "false":
                return "False"
            else:
                return repr(value)
        elif isinstance(value, list):
            items = [self.value_to_python_str(item) for item in value]
            return "[" + ", ".join(items) + "]"
        elif isinstance(value, dict):
            return self.dict_to_python_str(value)
        else:
            return repr(value)

def main():
    root = tk.Tk()
    app = YOLOTrainingGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()
